import os
import re
from io import StringIO
import numpy as np
from interpolate import Interpolator


class Grid:
    """
    实现对网格文件的读取解析
    """
    def __init__(self, grid_path: str, properties_path: str = None):
        """
        传入一个网格路径以及属性路径（如果有）
        :param grid_path: 网格路径
        :param properties_path: 属性值路径 非必须参数
        """
        self.grid_path = grid_path
        self.read(grid_path)
        self.properties = {}
        if properties_path is not None:
            properties = self.read_properties(properties_path)
            self.set_properties(properties)

    def set_properties(self, properties: dict):
        """
        设置属性值 传入一个字典
        :param properties: {"YOUNGMOD":np.array(...)}
        :return:
        """
        new_centroids = self.get_centroids()
        n_i, n_j, n_k = self.specgrid
        for key, prop in properties.items():
            # 这里将kji循环转成xyz对应的属性
            prop = prop.reshape(n_k, n_j, n_i)
            prop = np.transpose(prop, (2, 1, 0)).flatten()  # 确保最终顺序为 (x, y, z)
            self.properties[key] = Interpolator(new_centroids, prop)

    def reverse_convert_corner_point(self, corner: np.ndarray):
        """
        将焦点网格8哥顶点的数组 转换成坐标化顺序的数组
        convert_corner_point的逆向转换
        :param corner: (N,8)
        :return: array(Nz*2, Ny*2, Nx*2)
        """
        # 获取角点数组的形状，
        Nx, Ny, Nz = self.specgrid
        # 获取每层的角点数
        points_per_layer = Nx * 2 * Ny * 2
        # 构建 8 个角点的偏移量
        offsets = np.array([
            0, 1,
            Nx * 2, Nx * 2 + 1,  # 底部角点
            points_per_layer, points_per_layer + 1,  # 顶部角点前 2 个
            points_per_layer + Nx * 2, points_per_layer + Nx * 2 + 1  # 顶部角点后 2 个
        ])
        # 创建索引数组：每个单元的索引
        k_indices, j_indices, i_indices = np.indices((Nz, Ny, Nx))
        # 计算每个单元的基准索引
        k_base = k_indices * points_per_layer * 2 + j_indices * Nx * 2 * 2 + i_indices * 2
        # 将 corner 数据恢复到原始的 array 中
        array = np.zeros((Nz * 2, Ny * 2, Nx * 2))  # 创建一个新的数组，形状为 (Nz*2, Ny*2, Nx*2)
        # 扁平化角点数组（形状变为 (Nz * Ny * Nx, 8)）
        corner_flat = corner.reshape(-1, 8)
        # 生成每个格子对应的索引位置
        indices_flat = k_base.flatten()[:, None] + offsets
        # 使用 np.unravel_index 将索引转换为多维索引
        indices_multi = np.unravel_index(indices_flat, array.shape)
        # 将 corner_flat 中的数据放入 array 中
        array[indices_multi] = corner_flat
        # print(array.shape)
        return array

    def convert_corner_point(self, array: np.ndarray):
        """
        将坐标化的数组转成8顶点形式
        :param array: shape(N)
        :return: array(Nz, Ny, Nx, 8)
        """
        # 初始化每单元的  坐标数组，
        Nx, Ny, Nz = self.specgrid
        corner = np.zeros((Nz, Ny, Nx, 8), dtype=np.float64)
        # 每层的角点数
        points_per_layer = Nx * 2 * Ny * 2
        # 构建 8 个角点的偏移量
        offsets = np.array([
            0, 1,
            Nx * 2, Nx * 2 + 1,  # 底部角点
            points_per_layer, points_per_layer + 1,  # 顶部角点前 2 个
            points_per_layer + Nx * 2, points_per_layer + Nx * 2 + 1  # 顶部角点后 2 个
        ])
        # 创建索引数组：每个单元的索引
        k_indices, j_indices, i_indices = np.indices((Nz, Ny, Nx))
        # 计算每个单元的基准索引
        k_base = k_indices * points_per_layer * 2 + j_indices * Nx * 2 * 2 + i_indices * 2
        # 使用切片将 array 分配给 corner
        corner = corner.reshape(-1, 8)
        corner[:] = array[k_base.flatten()[:, None] + offsets]
        # 恢复原来的形状
        corner = corner.reshape(Nz, Ny, Nx, 8)
        return corner

    @staticmethod
    def parse_file(filename):
        """解析文件内容到字典"""
        with open(filename, 'r', encoding='utf8') as f:
            text = f.read()
        # 使用正则解析每个块
        blocks = re.split(r'(?<=/)\s*', text)
        data = {}
        for block in blocks:
            if not block.strip():
                continue
            lines = block.strip().split("\n")
            for i in range(len(lines)):
                key = lines.pop(0)
                if key.startswith("#"):
                    continue
                keyword = key.strip()
                break
            content = "\n".join(lines).strip(" /")
            data[keyword] = content
        return data

    @classmethod
    def read_properties(cls, filename):
        """读取所有属性并存入字典"""
        properties = {}
        if filename is not None and os.path.exists(filename):
            data = cls.parse_file(filename)
            for key, content in data.items():
                prop = cls.parse_propertie(content)
                properties[key] = prop
        return properties

    @staticmethod
    def parse_propertie(content):
        """基于 np.fromregex 的快速解析"""
        array = np.fromregex(
            StringIO(content),
            r'(\d+)\*([\d.]+)',
            dtype=[('num', np.int32), ('value', np.float64)]
        )
        return np.repeat(array['value'], array['num'])

    @staticmethod
    def read_specgrid(specgrid_text):
        """解析 SPECGRID 数据"""
        grid_list = re.findall(r'\d+', specgrid_text)[:3]
        if len(grid_list) < 3:
            raise ValueError("SPECGRID does not contain enough values for n_i, n_j, and n_k")
        n_i, n_j, n_k = map(int, grid_list)
        return [n_i, n_j, n_k]

    def read_coord(self, coord_text):
        """解析 COORD 数据"""
        n_i, n_j, n_k = self.specgrid
        array = np.fromstring(coord_text, sep=' ').reshape(-1, 6)
        coord_x = array[:, 0].reshape((n_j + 1, n_i + 1))
        coord_y = array[:, 1].reshape((n_j + 1, n_i + 1))
        # 根据ACTNUM 过滤下无效的单元格
        coord_x = self.get_valid_xy_coord(coord_x)
        coord_y = self.get_valid_xy_coord(coord_y)
        # 将最后有效的单元格维度存起来
        self.specgrid[0] = coord_x.shape[1] // 2
        self.specgrid[1] = coord_x.shape[0] // 2
        return coord_x, coord_y

    def get_valid_xy_coord(self, coord):
        """
        根据ACTNUM关键字的数组对xy的格子进行过滤
        返回有效的单元格数组
        :param coord:
        :return:
        """
        # 先扩充出重复格点
        A_expanded = np.repeat(coord, 2, axis=0)  # 沿着第一个维度（行）重复2次
        coord = np.repeat(A_expanded, 2, axis=1)  # 沿着第二个维度（列）重复2次
        coord = np.delete(coord, [0, -1], axis=0)
        coord = np.delete(coord, [0, -1], axis=1)
        # 将ACTNUM也扩充对应的的重复位点
        xy_actnum = self.actnum[0, :, :].repeat(2, 0).repeat(2, 1)
        coord = np.where(xy_actnum == 1, coord, np.nan)
        coord = coord[~np.isnan(coord).any(axis=1), :]  # 删除含 NaN 的行
        coord = coord[:, ~np.isnan(coord).any(axis=0)]  # 删除含 NaN 的列
        return coord

    def read_zcorn(self, zcorn_text):
        """解析 ZCORN 数据"""
        z = zcorn_text.split()
        z_data = np.array(z)
        zcorn = self.convert_corner_point(z_data)
        zactnum = self.actnum.astype(bool)
        z_data = zcorn[zactnum]
        return z_data

    def read(self, path):
        """读取文件并解析所有必要数据"""
        data = self.parse_file(path)
        self.specgrid = self.read_specgrid(data.pop("SPECGRID"))
        n_i, n_j, n_k = self.specgrid
        self.actnum = self.parse_propertie(data.pop("ACTNUM")).reshape(n_k, n_j, n_i)
        self.corner_z = self.read_zcorn(data.pop("ZCORN"))
        coord_x, coord_y = self.read_coord(data.pop("COORD"))
        # 重新获取下有效的格子维度
        n_i, n_j, n_k = self.specgrid
        # 将xy的coord 变成和z一样的 按照kji顺序的格子的8个顶点坐标 shape （N,8）
        self.corner_x = self.convert_corner_point(np.tile(coord_x, (n_k * 2, 1, 1)).flatten()).reshape(-1, 8)
        self.corner_y = self.convert_corner_point(np.tile(coord_y, (n_k * 2, 1, 1)).flatten()).reshape(-1, 8)
        self.other_data = data

    def get_centroids(self):
        """
        获取质心 用于插值计算
        :return:
        """
        centroids = np.stack((
            np.mean(self.corner_x, axis=1),
            np.mean(self.corner_y, axis=1),
            np.mean(self.corner_z, axis=1)
        ), axis=1)
        return centroids

    def convert_corner_to_grid_coordinates(self):
        """
        角点网格空间坐标化
        按照kji的循环顺序返回所有点的坐标
        :return: array(N,3)
        """
        x = self.reverse_convert_corner_point(self.corner_x)
        y = self.reverse_convert_corner_point(self.corner_y)
        z = self.reverse_convert_corner_point(self.corner_z)
        unique_corner_coords = np.column_stack((x.flatten(), y.flatten(), z.flatten()))
        return unique_corner_coords

    def convert_grid(self, new_grid_path):
        """
        根据新传入的路径 读取网格信息，并将现在网格的属性值 利用插值转换到新网格
        :param new_grid_path:
        :return: Grid()
        """
        grid = Grid(new_grid_path)
        new_centroids = grid.get_centroids()
        properties = {}
        for key, prop in self.properties.items():
            prop: Interpolator
            new_prop = prop.perform_interpolation(new_centroids)  # 对于目前题目的网格文件，已知为规则网格
            properties[key] = new_prop
        grid.set_properties(properties)
        return grid

    def write_keyword(self, fp, keyword, content, optimize=False):
        """
        根据关键词写入到grid.out
        :param fp: 文件句柄
        :param keyword: 关键字标签
        :param content: 内容 数组或者字符串
        :param optimize: 是否优化 即N*22323这种  用于输出属性
        :return:
        """
        fp.write(f"\n{keyword}\n")
        if isinstance(content, np.ndarray):
            if optimize:
                fp.write(self.format_array_optimized(content))
            else:
                np.savetxt(fp, content, fmt='%10.9f')
        else:
            fp.write(str(content))
        fp.write("\n")
        fp.write("/")

    @staticmethod
    def format_array_optimized(arr):
        """
        将arr转成N*1232323这种的字符串
        :param arr:
        :return:
        """
        # 计算相邻元素的差异
        diffs = np.diff(arr)
        # 找到变化的地方，包含第一个元素的起始位置
        change_points = np.where(diffs != 0)[0] + 1
        # 在变化的位置之间进行分段
        segments = np.split(arr, change_points)
        # 格式化每个段落
        result = [f"{len(segment)}*{segment[0]}" for segment in segments]
        # 每6个元素加一个换行符
        group_size = 6
        result_with_newlines = [' '.join(result[i:i + group_size]) for i in range(0, len(result), group_size)]
        return '\n'.join(result_with_newlines)

    def save_properties(self, properties_path="props.txt"):
        """
        输出属性信息
        :param properties_path:
        :return:
        """
        with open(properties_path, 'w', encoding='utf8') as f:
            for key in self.properties.keys():
                self.write_keyword(f, key, self[key], optimize=True)

    def __str__(self):
        return f"Grid({self.specgrid})"

    def __getitem__(self, item):
        if item in self.properties:
            # 现在是 x y z 转成 z y x
            n_i, n_j, n_k = self.specgrid
            property_array = self.properties[item].property_array
            property_array = property_array.reshape(n_i, n_j, n_k)
            property_array = np.transpose(property_array, (2, 1, 0))
            return property_array.flatten()
        raise KeyError()
